import numpy as np
import argparse
import cv2
from PIL import Image
from typing import Tuple, Dict, Optional
from skimage.metrics import peak_signal_noise_ratio
from quickMTF.quickMTF import quickMTF
import matplotlib.pyplot as plt


def compute_psnr(src_img: np.ndarray, ground_truth: Optional[np.ndarray] = None, is_plot=False):
    """
    计算图片的峰值信噪比，单位 dB
    参数:
    src_img: 输入图像 [H, W, C] or [H, W]
    ground_truth: 参考图像 [H, W, C] or [H, W, C], uint8 [0,255]。如果为None，则生成与img平均值相同的纯色buffer
    is_yuv: 输入图像是yuv格式吗？
    is_plot: 是否绘制图像
    返回:
    各通道峰值信噪比的平均值
    """
    channels = 1 if len(src_img.shape) == 2 else src_img.shape[-1]
    # 如果ground_truth为None，生成与img平均值相同的纯色buffer
    img = src_img
    gt = ground_truth
    if ground_truth is None:
        # 计算图像的平均值
        mean_value = np.mean(src_img, axis=(0, 1))
        # 生成与cropped_img相同大小的纯色buffer
        gt = np.full_like(src_img, mean_value)

    if channels == 1:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        gt = cv2.cvtColor(gt, cv2.COLOR_GRAY2RGB)
    # 计算每个通道的PSNR
    psnr_values = []
    for channel in range(channels):
        # 使用ground_truth作为参考图像
        psnr = peak_signal_noise_ratio(img[..., channel], gt[..., channel])
        psnr_values.append(psnr)
    
    if is_plot:
        plt.figure(figsize=(12, 6))
        plt.subplot(1, 2, 1)
        plt.title("Input Image")
        plt.imshow(img)
        plt.axis('off')
        plt.subplot(1, 2, 2)
        plt.title("Ground Truth")
        plt.imshow(gt)
        plt.axis('off')
        plt.tight_layout()
        plt.show()
    # 返回各通道PSNR的平均值
    return psnr_values[0:min(len(psnr_values), channels)]

def compute_mtf(img: np.ndarray, is_plot=False):
    """
    计算图片mtf。
    img: 输入单通道图像
    返回:
      mtf50和Nyquist frequency
    """
    
    # 创建quickMTF实例
    mtf_calculator = quickMTF()
    try:
        # 调用sfr_GUI方法计算MTF
        # 参数顺序: image, selected_index, counter, mtf_index/100, show_plots, return_fig
        mtf, angle, mtf_nyquist = mtf_calculator.sfr_GUI(
            img, 
            mtf_indx=0.5,
            show_plots = is_plot
        )
        return mtf, mtf_nyquist
    except Exception as e:
        print(f"MTF计算失败: {e}")
        return -1, -1

# -----------------------------
# main: 调用与测试
# -----------------------------
def main():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='图像质量评估工具')
    
    # 添加参数
    parser.add_argument('image_input', type=str, help='输入图像路径')
    parser.add_argument('-r', '--rect', type=int, nargs="+", required=True, 
                       help='选区坐标，格式: "x y width height"')
    parser.add_argument('-m', '--method', type=str, choices=['psnr', 'mtf'], help='计算方法: psnr(峰值信噪比), mtf(调制传递函数)')
    parser.add_argument('--is_plot', action='store_true', help='是否绘制图像')
    
    # 解析参数
    args = parser.parse_args()
    
    # 加载图像
    print(f"正在加载图像: {args.image_input}")
    image = Image.open(args.image_input)
    if image.mode in ['RGB', 'RGBA']:
        image = image.convert('RGB')
    image = np.array(image)
    x, y, width, height = args.rect
    # 确保选区在图像范围内
    img_height, img_width = image.shape[:2]
    x = max(0, x)
    y = max(0, y)
    width = min(width, img_width - x)
    height = min(height, img_height - y)
    
    # 裁剪选区
    cropped_img = image[y:y+height, x:x+width]

    # 显示参数信息
    print(f"\n参数信息:")
    print(f"  图像路径: {args.image_input}")
    print(f"  选区坐标: {args.rect}")
    print(f"  计算方法: {args.method}")
    print("-" * 50)
    # 根据method参数调用相应的计算方法
    if args.method == 'psnr':
        print("正在计算PSNR...")
        # 方式1: 不提供groundtruth，自动生成与图像平均值相同的纯色buffer
        cropped_yuv = cv2.cvtColor(cropped_img,  cv2.COLOR_RGB2YUV)
        psnr_values = compute_psnr(cropped_yuv, is_plot=args.is_plot)
        psnrs_str = ' '.join([f"{p:.2f}" for p in psnr_values])
        print(f"PSNR值(dB): ({psnrs_str})")
    
    elif args.method == 'mtf':
        print("正在计算MTF...")
        cropped_yuv = cv2.cvtColor(cropped_img,  cv2.COLOR_RGB2YUV)
        luma = cropped_yuv[:, :, 0]
        mtf_50, nyquist = compute_mtf(luma, is_plot=args.is_plot)
        if mtf_50 >= 0 and nyquist >= 0:
            print(f"MTF50: {mtf_50:.4f}")
            print(f"Nyquist MTF: {nyquist:.4f}")
        else:
            print("MTF计算失败")
    
    print("-" * 50)
    print("计算完成！")
        

if __name__ == "__main__":
    main()
